/**
 * Copyright (c) 2018-2022, NXOS Development Team
 * SPDX-License-Identifier: Apache-2.0
 *
 * Contains: Memory Manage Unite
 *
 * Change Logs:
 * Date           Author            Notes
 * 2022-1-16      JasonHu           Init
 * 2022-4-18      JasonHu           Add thead-c906 mmu support
 */
#include <base/types.h>
#include <base/defines.h>

#include <base/page.h>
#define NX_LOG_LEVEL NX_LOG_INFO
#include <base/log.h>
#include <platform.h>

#include <arch/mm.h>
#include <arch/pgtable_hwdef.h>
#include <arch/pgtable_prot.h>
#include <arch/pgtable_types.h>
#include <arch/pgtable.h>
#include <arch/base.h>
#include <base/memory.h>
#include <base/debug.h>
#include <base/irq.h>
#include <base/mmu.h>
#include <arch/sysregs.h>

#define NO_BLOCK_MAPPINGS BIT(0)

NX_IMPORT void HAl_EnableMmu(void);

#define GET_PF_ID(addr) ((addr) >> NX_PAGE_SHIFT)

NX_PRIVATE unsigned long AllocPageTable(void)
{
	void *phys;

	phys = NX_PageAlloc(1);
	if (phys != NX_NULL)
	{
		NX_MemZero(NX_Phy2Virt(phys), NX_PAGE_SIZE);
	}
	return (unsigned long)phys;
}

void PMD_SetSection(PMD_Type *pmdp, unsigned long phys,
					unsigned long prot)
{
	unsigned long sect_prot = PMD_TYPE_SECT | MakeSectionAttr(prot);

	PMD_Type new_pmd = PageMakePMD(phys >> PMD_SHIFT, sect_prot);

	NX_LOG_D("PMD_SetSection: %p prot:%p, pmd:%p", phys, prot, new_pmd);

	SetPMD(pmdp, new_pmd);
}

NX_PRIVATE void AllocInitPTE(PMD_Type *pmdp, unsigned long addr,
							 unsigned long end, unsigned long phys,
							 unsigned long prot,
							 unsigned long flags)
{
	PMD_Type pmd = *pmdp;
	PTE_Type *ptep;
	NX_Addr page;

	if (PMD_None(pmd))
	{
		unsigned long pte_phys;

		pte_phys = AllocPageTable();
		SetPMD(pmdp, MakePMD(pte_phys | PMD_TYPE_TABLE));
		pmd = *pmdp;
		page = NX_Virt2Phy((NX_Addr)pmdp & NX_PAGE_ADDR_MASK);
		NX_PageIncrease(page); /* increase pmd page */
	}

	ptep = PTE_OffsetVirt(pmdp, addr);
	page = NX_Virt2Phy((NX_Addr)ptep & NX_PAGE_ADDR_MASK);
	do
	{
		SetPTE(ptep, PageMakePTE(phys >> NX_PAGE_SHIFT, prot));
		NX_PageIncrease(page); /* increase pte page */
		phys += NX_PAGE_SIZE;
	} while (ptep++, addr += NX_PAGE_SIZE, addr != end);
}

NX_PRIVATE void AllocInitPMD(PUD_Type *pudp, unsigned long addr,
							 unsigned long end, unsigned long phys,
							 unsigned long prot,
							 unsigned long flags)
{
	PUD_Type pud = *pudp;
	PMD_Type *pmdp;
	unsigned long next;
	NX_Addr page;

	if (PUD_None(pud))
	{
		unsigned long pmd_phys;

		pmd_phys = AllocPageTable();
		SetPUD(pudp, MakePUD(pmd_phys | PUD_TYPE_TABLE));
		pud = *pudp;
		page = NX_Virt2Phy((NX_Addr)pudp & NX_PAGE_ADDR_MASK);
		NX_PageIncrease(page); /* increase pmd page */
	}

	pmdp = PMD_OffsetVirt(pudp, addr);
	do
	{
		next = PMD_AddrEnd(addr, end);

		if (((addr | next | phys) & ~SECTION_MASK) == 0 &&
			(flags & NO_BLOCK_MAPPINGS) == 0)
		{
			PMD_SetSection(pmdp, phys, prot);
			/* increase pmd page */
			page = NX_Virt2Phy((NX_Addr)pmdp & NX_PAGE_ADDR_MASK);
			NX_PageIncrease(page); /* increase pmd page */
		}
		else
			AllocInitPTE(pmdp, addr, next, phys,
						 prot, flags);

		phys += next - addr;
	} while (pmdp++, addr = next, addr != end);
}

NX_PRIVATE void AllocInitPUD(PGD_Type *pgdp, unsigned long addr,
							 unsigned long end, unsigned long phys,
							 unsigned long prot,
							 unsigned long flags)
{
	PGD_Type pgd = *pgdp;
	PUD_Type *pudp;
	unsigned long next;
	NX_Addr page;

	if (PGD_None(pgd))
	{
		unsigned long pud_phys;

		pud_phys = AllocPageTable();

		SetPGD(pgdp, MakePGD(pud_phys | PUD_TYPE_TABLE));
		pgd = *pgdp;
		page = NX_Virt2Phy((NX_Addr)pgdp & NX_PAGE_ADDR_MASK);
		NX_PageIncrease(page); /* increase pmd page */
	}

	pudp = PUD_OffsetVirt(pgdp, addr);
	do
	{
		next = PUD_AddrEnd(addr, end);
		AllocInitPMD(pudp, addr, next, phys,
					 prot, flags);
		phys += next - addr;

	} while (pudp++, addr = next, addr != end);
}

NX_PRIVATE void MapPGD(PGD_Type *pgdir, unsigned long phys,
					   unsigned long virt, unsigned long size,
					   unsigned long prot,
					   unsigned long flags)
{
	PGD_Type *pgdp = PGD_OffsetRaw(pgdir, virt);
	unsigned long addr, end, next;

	phys &= NX_PAGE_ADDR_MASK;
	addr = virt & NX_PAGE_ADDR_MASK;
	end = NX_PAGE_ALIGNUP(virt + size);

	do
	{
		next = PGD_AddrEnd(addr, end);
		AllocInitPUD(pgdp, addr, next, phys,
					 prot, flags);
		phys += next - addr;
	} while (pgdp++, addr = next, addr != end);
}

NX_PRIVATE void *__MapPageWithPhy_4KB(NX_Mmu *mmu, NX_Addr virAddr, NX_Addr phyAddr, NX_Size size, NX_UArch attr)
{
	/* attr -> prot */
	MapPGD(mmu->table, phyAddr, virAddr, size, attr, NO_BLOCK_MAPPINGS);
	return (void *)virAddr;
}

void *NX_HalMapPageWithPhy_4KB(NX_Mmu *mmu, NX_Addr virAddr, NX_Addr phyAddr, NX_Size size, NX_UArch attr)
{
	NX_ASSERT(mmu);
	if (!attr)
	{
		return NX_NULL;
	}

	virAddr = virAddr & NX_PAGE_ADDR_MASK;
	phyAddr = phyAddr & NX_PAGE_ADDR_MASK;
	size = NX_PAGE_ALIGNUP(size);

	NX_UArch level = NX_IRQ_SaveLevel();
	void *addr = __MapPageWithPhy_4KB(mmu, virAddr, phyAddr, size, attr);
	NX_IRQ_RestoreLevel(level);
	return addr;
}

NX_PRIVATE void *__MapPageWithPhy_2MB(NX_Mmu *mmu, NX_Addr virAddr, NX_Addr phyAddr, NX_Size size, NX_UArch attr)
{
	/* attr -> prot */
	MapPGD(mmu->table, phyAddr, virAddr, size, attr, 0);
	return (void *)virAddr;
}

void *NX_HalMapPageWithPhy_2MB(NX_Mmu *mmu, NX_Addr virAddr, NX_Addr phyAddr, NX_Size size, NX_UArch attr)
{
	NX_ASSERT(mmu);
	if (!attr)
	{
		return NX_NULL;
	}

	virAddr = virAddr & NX_PAGE_ADDR_MASK;
	phyAddr = phyAddr & NX_PAGE_ADDR_MASK;
	size = NX_PAGE_ALIGNUP(size);

	NX_UArch level = NX_IRQ_SaveLevel();
	void *addr = __MapPageWithPhy_2MB(mmu, virAddr, phyAddr, size, attr);
	NX_IRQ_RestoreLevel(level);
	return addr;
}

NX_PRIVATE unsigned long *walk_pte(PMD_Type *pmdp, unsigned long start, NX_Addr *entryTable[4])
{
	PTE_Type *ptep = PTE_OffsetVirt(pmdp, start);

	entryTable[0] = (NX_Addr *)ptep;

	NX_LOG_D("pte ref=%d", NX_PageReference(NX_Virt2Phy((NX_Addr)ptep & NX_PAGE_ADDR_MASK))); /* increase pmd page */

	return (unsigned long *)ptep;
}

NX_PRIVATE unsigned long *walk_pmd(PUD_Type *pudp, unsigned long start, NX_Addr *entryTable[4])
{
	unsigned long addr = start;
	PMD_Type *pmdp = PMD_OffsetVirt(pudp, start);
	PMD_Type pmd;

	pmd = *pmdp;

	entryTable[1] = (NX_Addr *)pmdp;

	NX_LOG_D("pmd ref=%d", NX_PageReference(NX_Virt2Phy((NX_Addr)pmdp & NX_PAGE_ADDR_MASK))); /* increase pmd page */

	if (PMD_None(pmd) || PMD_Section(pmd))
	{
		// NX_LOG_D("walk_pmd2: PMD_Section=%p, pmd=%p", PMD_Section(pmd), GetPMD(pmd) & SECTION_MASK);
		return (unsigned long *)pmdp;
	}
	else
		return walk_pte(pmdp, addr, entryTable);
}

NX_PRIVATE unsigned long *walk_pud(PGD_Type *pgdp, unsigned long start, NX_Addr *entryTable[4])
{
	unsigned long addr = start;
	PUD_Type *pudp = PUD_OffsetVirt(pgdp, start);
	PUD_Type pud;

	pud = *pudp;

	entryTable[2] = (NX_Addr *)pudp;

	NX_LOG_D("pud ref=%d", NX_PageReference(NX_Virt2Phy((NX_Addr)pudp & NX_PAGE_ADDR_MASK))); /* increase pmd page */

	if (PUD_None(pud) || PUD_Section(pud))
		return NX_NULL;
	else
		return walk_pmd(pudp, addr, entryTable);
}

NX_PRIVATE unsigned long *walk_pgd(PGD_Type *pgd, unsigned long start, NX_Addr *entryTable[4])
{
	unsigned long addr = start;
	PGD_Type *pgdp;
	PGD_Type pgd_entry;

	pgdp = PGD_OffsetRaw(pgd, start);

	if (!pgdp)
		return NX_NULL;
	pgd_entry = *pgdp;

	entryTable[3] = (NX_Addr *)pgdp;

	if (PGD_None(pgd_entry))
		return NX_NULL;
	else
		return walk_pud(pgdp, addr, entryTable);
}

NX_Addr *pgdir_walk(NX_Mmu *mmu, unsigned long vaddr, NX_Addr *entryTable[4])
{
	return walk_pgd(mmu->table, vaddr, entryTable);
}

NX_PRIVATE NX_Error UnmapOnePage(NX_Mmu *mmu, NX_Addr virAddr, NX_Bool freeLeaf)
{
	NX_Addr *entryTable[4];
	void *levelPageTable;
	NX_Bool freePage = NX_False;
	unsigned long *pmdp, *pudp, *pgdp;
	unsigned long *pte = pgdir_walk(mmu, virAddr & NX_PAGE_ADDR_MASK, entryTable);
	if (pte != NX_NULL)
	{
		PMD_Type pmd = *((PMD_Type *)pte);
		/**
		 * check pte is section or table
		 */
		if (PMD_None(pmd) || PMD_Section(pmd))
		{
			if (PMD_Section(pmd))
			{
				pmdp = pte;
				*pmdp = 0;
				freePage = NX_True;
			}
		}
		else
		{
			NX_Addr paddr = GetPMD(pmd) & PTE_ADDR_MASK;
			if (freeLeaf == NX_True)
			{
				NX_PageFree((void *)paddr);
			}
			*pte = 0;

			levelPageTable = (void *)(((NX_Addr)pte) & NX_PAGE_ADDR_MASK); /* get level page table by pte */
			/* 需要从pte开始检查释放 */
			NX_LOG_D("free pte %p page %p", pte, levelPageTable);
			NX_PageFree(levelPageTable);
			if (NX_PageReference(levelPageTable) == 1)
			{
				/* do real free pte */
				NX_PageFree(levelPageTable);

				pmdp = entryTable[1];
				*pmdp = 0;

				freePage = NX_True;
			}
		}
		if (freePage == NX_True)
		{
			/* 需要从pmd开始检查释放 */
			levelPageTable = (void *)(((NX_Addr)pmdp) & NX_PAGE_ADDR_MASK); /* get level page table by pte */
			NX_LOG_D("free pmd %p page %p", pmdp, levelPageTable);
			NX_PageFree(levelPageTable);
			if (NX_PageReference(levelPageTable) == 1)
			{
				/* do real free pmdp */
				NX_PageFree(levelPageTable);

				pudp = entryTable[2]; /* pud */
				*pudp = 0;
				levelPageTable = (void *)(((NX_Addr)pudp) & NX_PAGE_ADDR_MASK); /* get level page table by pte */
				NX_LOG_D("free pud %p page %p", pudp, levelPageTable);
				NX_PageFree(levelPageTable);
				if (NX_PageReference(levelPageTable) == 1)
				{
					/* do real free pudp */
					NX_PageFree(levelPageTable);

					pgdp = entryTable[3]; /* pgd */
					*pgdp = 0;
				}
			}
		}
	}
	return NX_EOK;
}

NX_INLINE NX_Error __UnmapPage(NX_Mmu *mmu, NX_Addr virAddr, NX_Size pages, NX_Bool freeLeaf)
{
	while (pages > 0)
	{
		UnmapOnePage(mmu, virAddr, freeLeaf);
		virAddr += NX_PAGE_SIZE;
		pages--;
	}
	return NX_EOK;
}

NX_PRIVATE NX_Error NX_HalUnmapPage(NX_Mmu *mmu, NX_Addr virAddr, NX_Size size, NX_Bool freeLeaf)
{
	NX_ASSERT(mmu);

	virAddr = virAddr & NX_PAGE_ADDR_MASK;
	size = NX_PAGE_ALIGNUP(size);

	NX_Addr addrStart = virAddr;
	NX_Addr addrEnd = virAddr + size - 1;
	NX_Size pages = GET_PF_ID(addrEnd) - GET_PF_ID(addrStart) + 1;

	NX_UArch level = NX_IRQ_SaveLevel();
	NX_Error err = __UnmapPage(mmu, virAddr, pages, freeLeaf);
	NX_IRQ_RestoreLevel(level);
	return err;
}

NX_PRIVATE void *NX_HalVir2Phy(NX_Mmu *mmu, NX_Addr virAddr)
{
	NX_ASSERT(mmu);

	NX_Addr paddr = 0;
	NX_Addr *entryTable[4];
	NX_Addr *pte = pgdir_walk(mmu, virAddr & NX_PAGE_ADDR_MASK, entryTable);
	if (pte != NX_NULL)
	{
		PMD_Type pmd = *((PMD_Type *)pte);
		if (PMD_None(pmd) || PMD_Section(pmd))
		{
			if (PMD_Section(pmd))
			{
				paddr = (GetPMD(pmd) & SECTION_MASK) & VA_MASK;
				paddr += virAddr & (SECTION_SIZE - 1);
			}
		}
		else
		{
			paddr = GetPMD(pmd) & PTE_ADDR_MASK;
			paddr += virAddr & (NX_PAGE_SIZE - 1);
		}
	}
	return (void *)paddr;
}

NX_PRIVATE NX_Bool IsVirAddrMapped(NX_Mmu *mmu, NX_Addr virAddr)
{
	NX_ASSERT(mmu);

	NX_Addr *entryTable[4];
	NX_Addr *pte = pgdir_walk(mmu, virAddr & NX_PAGE_ADDR_MASK, entryTable);
	if (pte != NX_NULL)
	{
		PMD_Type pmd = *((PMD_Type *)pte);
		if (PMD_None(pmd) || PMD_Section(pmd))
		{
			if (PMD_Section(pmd))
			{
				return NX_True;
			}
		}
		else
		{
			return NX_True;
		}
	}
	return NX_False;
}

NX_PRIVATE NX_Error MapOnePage(NX_Mmu *mmu, NX_Addr virAddr, NX_Addr phyAddr, NX_UArch attr)
{
	if (IsVirAddrMapped(mmu, virAddr) == NX_True)
	{
		NX_LOG_E("map page: vir:%p was mapped!", virAddr);
		return NX_EINVAL;
	}

	MapPGD(mmu->table, phyAddr, virAddr, NX_PAGE_SIZE, attr, NO_BLOCK_MAPPINGS);

	return NX_EOK;
}

NX_PRIVATE void *__MapPage(NX_Mmu *mmu, NX_Addr virAddr, NX_Size size, NX_UArch attr)
{
	NX_Addr addrStart = virAddr;
	NX_Addr addrEnd = virAddr + size - 1;

	NX_SSize pages = GET_PF_ID(addrEnd) - GET_PF_ID(addrStart) + 1;
	NX_Size mappedPages = 0;
	void *phyAddr;

	while (pages > 0)
	{
		phyAddr = NX_PageAlloc(1);
		if (phyAddr == NX_NULL)
		{
			NX_LOG_E("map page: alloc page failed!");
			goto err;
		}

		if (MapOnePage(mmu, virAddr, (NX_Addr)phyAddr, attr) != NX_EOK)
		{
			NX_LOG_E("map page: vir:%p phy:%p attr:%x failed!", virAddr, phyAddr, attr);
			goto err;
		}
		virAddr += NX_PAGE_SIZE;
		phyAddr += NX_PAGE_SIZE;
		pages--;
		mappedPages++;
	}
	return (void *)addrStart;
err:
	if (phyAddr != NX_NULL)
	{
		NX_PageFree(phyAddr);
	}
	__UnmapPage(mmu, addrStart, mappedPages, NX_True);
	return NX_NULL;
}

NX_PRIVATE void *NX_HalMapPage(NX_Mmu *mmu, NX_Addr virAddr, NX_Size size, NX_UArch attr)
{
	NX_ASSERT(mmu);
	if (!attr)
	{
		return NX_NULL;
	}

	virAddr = virAddr & NX_PAGE_ADDR_MASK;
	size = NX_PAGE_ALIGNUP(size);

	NX_UArch level = NX_IRQ_SaveLevel();
	void *addr = __MapPage(mmu, virAddr, size, attr);
	NX_IRQ_RestoreLevel(level);
	return addr;
}

NX_PRIVATE void *NX_HalMapPageWithPhy(NX_Mmu *mmu, NX_Addr virAddr, NX_Addr phyAddr, NX_Size size, NX_UArch attr)
{
	return NX_HalMapPageWithPhy_4KB(mmu, virAddr, phyAddr, size, attr);
}

#define TTBR0_EL12 "S3_5_C2_C0_0"

#define GET_SYS_REG(reg, out) __asm__ volatile("mrs %0, " reg \
											   : "=r"(out)::"memory");
#define SET_SYS_REG(reg, in) __asm__ volatile("msr " reg ", %0" ::"r"(in) \
											  : "memory");

NX_PRIVATE void NX_HalSetPageTable(NX_Addr addr)
{
	write_sysreg(addr, ttbr0_el1);
	// NX_LOG_I("NX_HalSetPageTable");
	// SET_SYS_REG(TTBR0_EL2, addr);
	// NX_LOG_I(" SetPageTable done ");

}

NX_PRIVATE NX_Addr NX_HalGetPageTable(void)
{
	NX_Addr reg = read_sysreg(ttbr0_el1);
	// NX_Addr reg;
	// NX_LOG_I("NX_HalGetPageTable");

	// GET_SYS_REG(TTBR0_EL12, reg);
	// NX_LOG_I(" GetPageTable done ");

	return reg;
}

NX_PRIVATE void NX_HalEnable(void)
{
	HAl_EnableMmu();
}

NX_PRIVATE void NX_HalFlush(void)
{
	dsb(ishst);
	asm("tlbi	vmalle1is");
	dsb(ish);
	isb();
}

NX_IMPORT void DumpPageTable(PGD_Type *pgd, NX_Addr virAddr, NX_Size size);
NX_PRIVATE void NX_HalDump(NX_Mmu *mmu, NX_Addr virAddr, NX_Size size)
{
	DumpPageTable((PGD_Type *)mmu->table, virAddr, size);
}

NX_INTERFACE struct NX_MmuOps NX_MmuOpsInterface =
	{
		.setPageTable = NX_HalSetPageTable,
		.getPageTable = NX_HalGetPageTable,
		.enable = NX_HalEnable,
		.mapPage = NX_HalMapPage,
		.mapPageWithPhy = NX_HalMapPageWithPhy,
		.unmapPage = NX_HalUnmapPage,
		.vir2Phy = NX_HalVir2Phy,
		.flush = NX_HalFlush,
		.dump = NX_HalDump,
};
